import psycopg
from psycopg_pool import AsyncConnectionPool
from common.logger import logger
from contextlib import asynccontextmanager
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.store.postgres.aio import AsyncPostgresStore


class PostgresManager:
    def __init__(self, db_uri: str, min_size: int = 1, max_size: int = 20):
        self.db_uri = db_uri
        self.min_size = min_size
        self.max_size = max_size
        self.pool: AsyncConnectionPool | None = None

    @asynccontextmanager
    async def get_pool(self):
        """使用上下文管理器方式获取连接池"""
        async with AsyncConnectionPool(
            conninfo=self.db_uri,
            min_size=self.min_size,
            max_size=self.max_size,
        ) as pool:
            yield pool

    async def initialize(self) -> AsyncConnectionPool:
        """初始化连接池"""
        logger.info("调用 PostgresManager.initialize 函数")
        if not self.pool:
            # 使用推荐的方式创建连接池
            self.pool = AsyncConnectionPool(
                conninfo=self.db_uri,
                min_size=self.min_size,
                max_size=self.max_size,
                open=False
            )
            await self.pool.open()
        return self.pool

    async def close(self):
        """关闭连接池"""
        logger.info("调用 PostgresManager.close 函数")
        if self.pool:
            await self.pool.close()
            self.pool = None

    async def init_checkpointer_store_schema(self, init_start: bool = False):
        """初始化 checkpointer schema"""
        if init_start:
            logger.info("初始化 Postgres checkpointer&store schema")
            # 使用独立连接执行 schema 初始化，避免事务冲突
            conn = await psycopg.AsyncConnection.connect(
                self.db_uri, autocommit=True
            )
            try:
                temp_checkpointer = AsyncPostgresSaver(conn)
                await temp_checkpointer.setup()
                temp_store = AsyncPostgresStore(conn)
                await temp_store.setup()
                logger.info("✅ Checkpointer&Store schema 初始化完成")
            finally:
                await conn.close()

    def get_checkpointer(self, pool: AsyncConnectionPool) -> AsyncPostgresSaver:
        """获取 checkpointer 实例"""
        return AsyncPostgresSaver(pool)

    def get_store(self, pool: AsyncConnectionPool) -> AsyncPostgresStore:
        """获取 store 实例"""
        return AsyncPostgresStore(pool)


@asynccontextmanager
async def postgres_lifespan(db_uri: str, init_start: bool = False, min_size: int = 1, max_size: int = 20):
    """PostgreSQL 生命周期管理器"""
    # 使用推荐的上下文管理器方式管理 Postgres 连接池
    async with AsyncConnectionPool(
        conninfo=db_uri,
        min_size=min_size,
        max_size=max_size,
    ) as pool:
        # 创建 PostgresManager 实例
        postgres_manager = PostgresManager(db_uri, min_size, max_size)
        
        # 初始化 checkpointer schema（仅启动时执行一次）
        await postgres_manager.init_checkpointer_store_schema(init_start)
        
        # 创建 checkpointer 实例
        checkpointer = postgres_manager.get_checkpointer(pool)
        store = postgres_manager.get_store(pool)
        
        yield pool, checkpointer, store
        
        logger.info("🛑 Postgres 连接池自动关闭")
